package com.atguigu.study.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * @author zzyy
 * @create 2025-03-03 12:35
 */
@RestController
@Slf4j
public class EmbeddinglController {
    @Resource
    private EmbeddingModel embeddingModel; // 文本向量化模型
    @Resource
    private QdrantClient qdrantClient;// 向量数据库访问的连接客户端
    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;// 对向量数据库CRUD的操作类

    /**
     * 文本向量化测试，看看形成向量后的文本
     * http://localhost:9010/embedding/embed
     *
     * @return
     */
    @GetMapping(value = "/embedding/embed")
    public String embed() {
        String prompt = """
               咏鸡
            鸡鸣破晓光，
            红冠映朝阳。
            金羽披霞彩，
            昂首步高岗。
            """;
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);

        System.out.println(embeddingResponse);

        return embeddingResponse.content().toString();
    }

    /**
     * 新建向量数据库实例和创建索引：test-qdrant
     * 类似mysql create database test-qdrant
     * <p>
     * http://localhost:9010/embedding/createCollection
     */
    @GetMapping(value = "/embedding/createCollection")
    public void createCollection() {
        var vectorParams = Collections.VectorParams.newBuilder()
            .setDistance(Collections.Distance.Cosine)
            .setSize(1024)
            .build();
        qdrantClient.createCollectionAsync("test-qdrant", vectorParams);
    }

    /**
     * 往向量数据库新增文本记录
     * http://localhost:9010/embedding/add
     */
    @GetMapping(value = "/embedding/add")
    public String add() {
        String prompt = """
            咏鸡
            鸡鸣破晓光，
            红冠映朝阳。
            金羽披霞彩，
            昂首步高岗。
            """;
        TextSegment segment1 = TextSegment.from(prompt);
        segment1.metadata().put("author", "Beerus");
        segment1.metadata().put("UserId", "1111");
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        String result = embeddingStore.add(embedding1, segment1);

        System.out.println(result);

        return result;
    }

    @GetMapping(value = "/embedding/query1")
    public void query1() {
        Embedding queryEmbedding = embeddingModel.embed("咏鸡现代诗").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
            .queryEmbedding(queryEmbedding)
            .maxResults(1)
            .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        System.out.println(searchResult.matches().get(0).embedded().text());
    }

    @GetMapping(value = "/embedding/query2")
    public void query2() {
        Embedding queryEmbedding = embeddingModel.embed("咏鸡").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
            .queryEmbedding(queryEmbedding)
            .filter(metadataKey("author").isEqualTo("zzyy"))
            .maxResults(1)
            .build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
    }
}
